import tqdm
from gxl_ai_utils.thread.my_thread import GxlDynamicThreadPool
from gxl_ai_utils.utils import utils_file, utils_data


def get_tsv_from_wav_scp(wav_scp_path, output_tsv_path, num_thread=32):
    """"""
    utils_file.logging_print(f"get_tsv_from_wav_scp:开始处理如下wav_scp: {wav_scp_path if isinstance(wav_scp_path, str) else "dict不予展示"}, tsv_path: {output_tsv_path}")
    utils_file.makedir_sil(output_tsv_path)
    if isinstance(wav_scp_path, str):
        wav_dict = utils_file.load_dict_from_scp(wav_scp_path)
    elif isinstance(wav_scp_path, dict):
        wav_dict = wav_scp_path
    else:
        raise ValueError("get_tsv_from_wav_scp: wav_scp_path must be str or dict")
    res_list = ["/"]
    wav_list = list(wav_dict.values())
    list_list = utils_file.do_split_list(wav_list, num_thread)
    runner = GxlDynamicThreadPool()
    for list_i in list_list:
        """"""
        runner.add_task(little_fun4get_tsv_from_wav_scp, [res_list, list_i])
    runner.start()
    utils_file.write_list_to_file(res_list, "./all_data_with_sample.tsv")


def little_fun4get_tsv_from_wav_scp(res_list, wav_path_list):
    temp_list = []
    for wav_path in tqdm.tqdm(wav_path_list, total=len(wav_path_list)):
        """"""
        samples, _ = utils_data.get_sample_count(wav_path)
        temp_list.append(f"{wav_path}\t{samples}")
    res_list.extend(temp_list)




def prepare_data():
    """"""
    train_wav_path = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_aishell/examples/aishell/s0/data/train/wav.scp"
    dev_wav_path = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_aishell/examples/aishell/s0/data/dev/wav.scp"
    get_tsv_from_wav_scp(train_wav_path, "./data/train_aishell.tsv")
    get_tsv_from_wav_scp(dev_wav_path, "./data/dev_aishell.tsv")

if __name__ == "__main__":
    prepare_data()